# GDP与用电量关联分析散点图.py
# 功能：读取GDP和用电量累计值，计算同比增速，绘制散点图（X轴：用电量同比，Y轴：GDP同比，下拉切换省份）

import os
import pandas as pd
import plotly.graph_objects as go
from plotly.subplots import make_subplots

# ------------------ 路径设置 ------------------
abs_path = os.path.abspath('.')
data_dir = os.path.join(f'{abs_path}/关联分析', 'data')
gdp_path = os.path.join(data_dir, 'GDP累计值.xlsx')
power_path = os.path.join(data_dir, '用电量累计值.xlsx')
output_html = os.path.join(f'{abs_path}/output', 'GDP与用电量同比增速散点图分析.html')

# ------------------ 读取数据 ------------------
# GDP：跳过第一行，第二行为列名，第一列为日期
gdp_df = pd.read_excel(gdp_path, header=1, index_col=0)
gdp_df.index = pd.to_datetime(gdp_df.index)

# 用电量：第一行为列名，第一列为日期
power_df = pd.read_excel(power_path, header=0, index_col=0)
power_df.index = pd.to_datetime(power_df.index)

# ------------------ 提取省份 ------------------
def extract_province(col):
    return col.split(':')[0]

gdp_cols = {extract_province(col): col for col in gdp_df.columns}
power_cols = {extract_province(col): col for col in power_df.columns}

# 共有省份
common_provinces = sorted(set(gdp_cols.keys()) & set(power_cols.keys()))

if len(common_provinces) == 0:
    raise ValueError("没有共同的省市数据，请检查列名格式是否为 '省份:指标:类型'")

print(f"共找到 {len(common_provinces)} 个共有省市：{common_provinces}")

# ------------------ 计算同比函数 ------------------
def calculate_yoy(series):
    """计算同比增速：(当前值 - 去年同期值) / |去年同期值|，返回单位为%的Series"""
    series = series.dropna()
    yoy = []
    dates = []
    for date, value in series.items():
        if date.month in [3, 6, 9, 12]:  # 只处理季度数据
            last_year_date = date - pd.DateOffset(years=1)
            if last_year_date in series.index:
                last_year_value = series[last_year_date]
                if abs(last_year_value) > 0:  # 避免除以0
                    growth = (value - last_year_value) / abs(last_year_value) * 100
                    yoy.append(growth)
                    dates.append(date)
    return pd.Series(yoy, index=dates).sort_index()

# ------------------ 创建图表 ------------------
fig = go.Figure()

# 默认显示第一个省份
first_prov = common_provinces[0]
gdp_col = gdp_cols[first_prov]
power_col = power_cols[first_prov]

# 计算同比并合并数据（按日期对齐）
gdp_yoy = calculate_yoy(gdp_df[gdp_col].dropna())
power_yoy = calculate_yoy(power_df[power_col].dropna())
merged_data = pd.concat([power_yoy, gdp_yoy], axis=1, keys=['power', 'gdp']).dropna()

# 添加初始散点
fig.add_trace(
    go.Scatter(
        x=merged_data['power'],
        y=merged_data['gdp'],
        mode='lines+markers',
        name=f'{first_prov} 数据点',
        marker=dict(
            size=8,
            color='blue',
            opacity=0.7
        ),
        text=[f"日期: {idx.strftime('%Y-%m')}<br>用电量同比: {x:.2f}%<br>GDP同比: {y:.2f}%" 
              for idx, x, y in zip(merged_data.index, merged_data['power'], merged_data['gdp'])],
        hoverinfo='text'
    )
)

# ------------------ 构建下拉菜单 ------------------
dropdown_buttons = []

for province in common_provinces:
    gdp_col = gdp_cols[province]
    power_col = power_cols[province]

    # 计算同比并合并数据
    gdp_yoy = calculate_yoy(gdp_df[gdp_col].dropna())
    power_yoy = calculate_yoy(power_df[power_col].dropna())
    merged_data = pd.concat([power_yoy, gdp_yoy], axis=1, keys=['power', 'gdp']).dropna()

    # 构建下拉按钮配置
    button = dict(
        label=province,
        method='update',
        args=[
            {
                'x': [merged_data['power']],
                'y': [merged_data['gdp']],
                'text': [f"日期: {idx.strftime('%Y-%m')}<br>用电量同比: {x:.2f}%<br>GDP同比: {y:.2f}%" 
                         for idx, x, y in zip(merged_data.index, merged_data['power'], merged_data['gdp'])],
                'name': f'{province} 数据点'
            },
            {'title': f"{province} GDP同比增速 vs 用电量同比增速"}
        ]
    )
    dropdown_buttons.append(button)

# ------------------ 布局设置 ------------------
fig.update_layout(
    updatemenus=[
        dict(
            buttons=dropdown_buttons,
            direction='down',
            showactive=True,
            x=0.1,
            xanchor='left',
            y=1.15,
            yanchor='top',
            font=dict(size=12),
            bordercolor="#ccc",
            borderwidth=1,
            bgcolor="#f9f9f9"
        )
    ],
    title={
        'text': f'各省市GDP同比增速 vs 用电量同比增速（当前：{first_prov}）',
        'x': 0.5,
        'xanchor': 'center',
        'y': 0.95,
        'font': {'size': 18}
    },
    xaxis_title="用电量同比增速（%）",
    yaxis_title="GDP同比增速（%）",
    hovermode="closest",
    height=700,
    margin=dict(t=150),
    xaxis=dict(
        zeroline=True,
        zerolinecolor='gray',
        zerolinewidth=1
    ),
    yaxis=dict(
        zeroline=True,
        zerolinecolor='gray',
        zerolinewidth=1
    )
)

# ------------------ 输出HTML ------------------
fig.write_html(
    output_html,
    config={'displayModeBar': True},
    include_plotlyjs='cdn',
    auto_open=False
)

print(f"✅ 已生成散点图，保存至：{output_html}")
print("💡 用浏览器打开，使用顶部下拉菜单切换省份。")
print("📈 X轴：用电量同比增速，Y轴：GDP同比增速")